/**
 * @file uring_client.cc
 * @author SHADOW (shadow_yuan@qq.com)
 * @brief tcp client demo (io_uring)
 * @version 0.1
 * @date 2024-10-11
 *
 */
#ifdef _WIN32
#include <WinSock2.h>
#include <WS2tcpip.h>
#include <MSWSock.h>
#else
#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <liburing.h>
#endif // _WIN32
#include <string.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <errno.h>

#include <atomic>
#include <chrono>
#include <string>
#include <iostream>
#include <memory>
#include <deque>
#include <thread>

#ifndef _WIN32
typedef int SOCKET;
#define CLOSESOCKET close
#else
#define CLOSESOCKET closesocket
#pragma comment(lib, "ws2_32.lib")
#endif // !_WIN32

typedef std::atomic<intptr_t> SpinLock;

constexpr int kNullEvent = 0;
constexpr int kAcceptEvent = 1;
constexpr int kConnectEvent = 1 << 1;
constexpr int kSendEvent = 1 << 2;
constexpr int kRecvEvent = 1 << 3;

struct DataNode {
    int event;
    int size;
    char data[];
};

constexpr int kEmptyNodeSize = sizeof(DataNode);
constexpr int kMaxBuffer = 1024;

#define DataNodePtr(node) (&node->data[0])

static inline DataNode *MakeDataNode(int size) {
    char *buff = new char[kEmptyNodeSize + size];
    DataNode *node = reinterpret_cast<DataNode *>(buff);
    node->event = kNullEvent;
    node->size = size;
    return node;
}

static inline DataNode *MakeDataNode(int size, int event) {
    char *buff = new char[kEmptyNodeSize + size];
    DataNode *node = reinterpret_cast<DataNode *>(buff);
    node->event = event;
    node->size = size;
    return node;
}

static inline void FreeDataNode(DataNode *node) { delete[] (char *)node; }

static inline bool spinlock_trylock(SpinLock *lock) {
    intptr_t expected = 0;
    return lock->compare_exchange_strong(expected, 1, std::memory_order_acquire,
                                         std::memory_order_relaxed);
}

static inline void spinlock_unlock(SpinLock *lock) { lock->store(0, std::memory_order_release); }

#define spinlock_lock(lock)                                                                        \
    do {                                                                                           \
    } while (!spinlock_trylock(lock))

int set_socket_nonblocking(SOCKET fd, int nonblocking) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags < 0) return -1;
    if (nonblocking) {
        flags |= O_NONBLOCK;
    } else {
        flags &= ~O_NONBLOCK;
    }
    return fcntl(fd, F_SETFL, flags);
}

SpinLock g_lock(0);
bool g_send_pending = false;
bool g_connected = false;
struct io_uring g_ring;
std::deque<DataNode *> g_send_queue;
bool g_running = true;

void AsyncRecv(SOCKET fd) {
    auto node = MakeDataNode(kMaxBuffer, kRecvEvent);
    struct io_uring_sqe *sqe = io_uring_get_sqe(&g_ring);
    io_uring_prep_recv(sqe, fd, DataNodePtr(node), node->size, 0);
    io_uring_sqe_set_data(sqe, node);
    io_uring_submit(&g_ring);
}

void AsyncSend(SOCKET fd) {
    if (g_send_queue.empty() || g_send_pending) {
        return;
    }

    g_send_pending = true;

    DataNode *node = NULL;
    spinlock_lock(&g_lock);
    node = g_send_queue.front();
    g_send_queue.pop_front();
    spinlock_unlock(&g_lock);
    g_send_pending = true;

    struct io_uring_sqe *sqe = io_uring_get_sqe(&g_ring);
    io_uring_prep_send(sqe, fd, DataNodePtr(node), node->size, 0);
    io_uring_sqe_set_data(sqe, node);
    io_uring_submit(&g_ring);
}

void thread_proc(SOCKET sockfd) {
    SOCKET tmp = 0;
    while (g_running) {
        io_uring_cqe *cqe = NULL;
        int ret = io_uring_wait_cqe(&g_ring, &cqe);
        if (ret != 0) {
            printf("io_uring_wait_cqe return %d, cqe: %p\n", ret, cqe);
            if (cqe) {
                DataNode *node = (DataNode *)io_uring_cqe_get_data(cqe);
                io_uring_cqe_seen(&g_ring, cqe);
                FreeDataNode(node);
            }
            continue;
        }

        DataNode *node = (DataNode *)io_uring_cqe_get_data(cqe);
        switch (node->event) {
        case kConnectEvent:
            tmp = *(SOCKET *)DataNodePtr(node);
            if (cqe->res == 0) {
                printf("connect event, [success] fd: %d\n", tmp);
                g_connected = true;
                AsyncRecv(tmp);
            } else {
                printf("connect event, [failure] fd: %d res: %d\n", tmp, cqe->res);
                g_connected = false;
            }
            break;
        case kSendEvent:
            g_send_pending = false;
            printf("send event, written %d bytes\n", cqe->res);
            AsyncSend(sockfd);
            break;
        case kRecvEvent:
            if (cqe->res == 0) {
                printf("recv event, connection closed!\n");
                g_connected = false;
            } else {
                printf("recv event, recv %d bytes : '%s'\n", cqe->res, DataNodePtr(node));
                AsyncRecv(sockfd);
            }
            break;
        }
        io_uring_cqe_seen(&g_ring, cqe);
        // io_uring_peek_cqe(&g_ring, &cqe);
        FreeDataNode(node);
    }
}

void str_cli(SOCKET sockfd) {
    char sendbuf[kMaxBuffer];

    while (fgets(sendbuf, kMaxBuffer, stdin) != NULL) {
        int bytes = (int)strlen(sendbuf);
        while (bytes > 0) {
            if (sendbuf[bytes - 1] != '\r' && sendbuf[bytes - 1] != '\n') {
                break;
            }
            bytes -= 1;
        }

        if (bytes > 0) {
            auto node = MakeDataNode(bytes, kSendEvent);
            memcpy(DataNodePtr(node), sendbuf, bytes);

            spinlock_lock(&g_lock);
            g_send_queue.push_back(node);
            spinlock_unlock(&g_lock);

            if (g_connected) AsyncSend(sockfd);
        }
    }
}

int main(int argc, char *argv[]) {
    struct sockaddr_in server_addr;
    int len, port, i;
    SOCKET sockfd;
    char *address;
    char ip_str[256];
    std::thread thd;

    if (argc != 2) {
        printf("usage: client <ip:port>\n");
        return -1;
    }

#ifdef _WIN32
    WSADATA data;
    int status;
    status = WSAStartup(MAKEWORD(2, 0), &data);
    if (status != 0) {
        printf("failed to init win sock\n");
        return -1;
    }
#endif // _WIN32

    address = argv[1];
    len = (int)strlen(address);

    for (i = 0; i < len; i++) {
        if (address[i] == ':') {
            goto done;
        }
    }
    printf("usage: input param format error\n");
    return -1;
done:
    // init io_uring object
    io_uring_queue_init(64, &g_ring, 0);

    memcpy(ip_str, address, i);
    ip_str[i] = 0;

    port = atoi(&address[i + 1]);

    printf("prepare to connect server '%s:%d'\n", ip_str, port);

    if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) == (SOCKET)(-1)) {
        printf("failed to create socket\n");
        return -1;
    }

    set_socket_nonblocking(sockfd, 1);

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);

    if (inet_pton(AF_INET, ip_str, &server_addr.sin_addr) < 0) {
        printf("inet_pton error for: %s\n", ip_str);
        return -1;
    }

    thd = std::thread(thread_proc, sockfd);
    // r = connect(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr));
    auto sqe = io_uring_get_sqe(&g_ring);
    io_uring_prep_connect(sqe, sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr));
    auto node = MakeDataNode(sizeof(sockfd), kConnectEvent);
    *(SOCKET *)DataNodePtr(node) = sockfd;
    io_uring_sqe_set_data(sqe, node);
    io_uring_submit(&g_ring);
    printf("submit async connect, fd: %d\n", sockfd);

    str_cli(sockfd);
    CLOSESOCKET(sockfd);
#ifdef _WIN32
    WSACleanup();
#endif // _WIN32
    return 0;
}
